import os

from langchain.agents import AgentType
from langchain.agents import initialize_agent
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain.prompts import PromptTemplate
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.tools.retriever import create_retriever_tool
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain.agents import AgentExecutor

# 构造TAVILY搜索工具
os.environ["TAVILY_API_KEY"] = "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"
search = TavilySearchResults(max_results=2)
# result = search.invoke("what is the weather in SF")
# print(result)

# 构造向量查询工具
embeddings = OpenAIEmbeddings(base_url="https://api.openai-hk.com/v1",
                              api_key="hk-0amgwp10000255022bdd816341db25b54dc2e46787aee69f")
loader = WebBaseLoader("https://docs.smith.langchain.com/overview")
docs = loader.load()
documents = RecursiveCharacterTextSplitter(
    chunk_size=1000, chunk_overlap=200
).split_documents(docs)
vector = FAISS.from_documents(documents, embeddings)
retriever = vector.as_retriever()
# print(retriever.invoke("how to upload a dataset")[0])
retriever_tool = create_retriever_tool(
    retriever,
    "langsmith_search",
    "Search for information about LangSmith. For any questions about LangSmith, you must use this tool!",
)

# 封装大模型可以使用的工具包
tools = [search, retriever_tool]

# 构造大模型并绑定工具
model = ChatOpenAI(base_url="https://api.openai-hk.com/v1",
                   api_key="hk-0amgwp10000255022bdd816341db25b54dc2e46787aee69f")
prompt = ChatPromptTemplate.from_messages(
    [SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='You are a helpful assistant')),
     MessagesPlaceholder(variable_name='chat_history', optional=True),
     HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], template='{input}')),
     MessagesPlaceholder(variable_name='agent_scratchpad')])
agent = initialize_agent(tools, model, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, prompt=prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools)

response = agent.invoke({"input": "hi!", "page_content": "how to upload a dataset"})
print(response)
